
import pdb
import math
import cv2
import numpy as np
import sys
import numpy as np
import os.path as osp
import cv2
import sys
def showImg(img,wati_time=0,win_name="r"):
    print(img.shape)
    cv2.namedWindow(win_name,cv2.WINDOW_NORMAL)
    cv2.imshow(win_name,img)
    key = cv2.waitKey(wati_time) & 0xFF
    if key == ord('q'):
        sys.exit()
    return key

def getKD(K_p):
  fs = cv2.FileStorage(K_p, cv2.FILE_STORAGE_READ)
  k_node = fs.getNode("K")
  K = k_node.mat()
  d_node = fs.getNode("D")
  D = d_node.mat()
  print("K: \n",K)
  print("D: \n",D)
  return K,D


def select_points(image):
    points = []
    win_name='Select Points'
    cv2.namedWindow(win_name,cv2.WINDOW_NORMAL)

    def mouse_callback(event, x, y, flags, param):
        nonlocal points
        if event == cv2.EVENT_RBUTTONDOWN:
            points.append((x, y))
            cv2.circle(image, (x, y), 2, (0, 255, 0), -1)
            cv2.imshow(win_name, image)
    cv2.setMouseCallback(win_name, mouse_callback)

    while True:
        cv2.imshow(win_name, image)
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'): 
            break

    cv2.destroyAllWindows()

    return np.array(points,dtype=np.float)


def get_Eb2c(K,D, raw_img, board, dictionary,debug=True):
    arucoParams = cv2.aruco.DetectorParameters()
    detector = cv2.aruco.ArucoDetector(dictionary, arucoParams)

    frame_remapped_gray = cv2.cvtColor(raw_img, cv2.COLOR_BGR2GRAY)
    corners, ids, aruco_rejected = detector.detectMarkers(frame_remapped_gray,)
    corners, ids, rejectedImgPoints, _ = detector.refineDetectedMarkers(
                                                frame_remapped_gray, 
                                                board, 
                                                corners, 
                                                ids, 
                                                rejectedCorners=aruco_rejected, 
                                                cameraMatrix=K, 
                                                distCoeffs=D)
    # print("**********: rejectedImgPoints: \n",rejectedImgPoints)
    # print("**********: ids: \n",ids)
    if len(corners) == 0:
        print("ERROR: can not find aruco ID pts \n")
        exit()
    im_with_aruco_board = cv2.aruco.drawDetectedMarkers(raw_img.copy(), corners, ids, (0,255,0))
    obj_points, img_points = board.matchImagePoints(corners, ids,ids)

    save_corners_txt = False
    if save_corners_txt:
        save_obs_img_pts(".",img_points,obj_points,ids)

    # print("ids \n", ids)
    # print("obj_points \n", obj_points)
    # print("len obj_points", len(obj_points))

    # print("img_points", img_points)
    # print("len img_points", len(img_points))
    ok, rvec, tvec = cv2.solvePnP(obj_points, img_points, K, D)

    show_corners = True
    if show_corners:
        for pts in img_points:
            # print(pts)
            cv2.circle(im_with_aruco_board, (int(pts[0][0]),int(pts[0][1])), 5, [0, 255, 0],-1)
        img_undis=cv2.undistort(im_with_aruco_board,K,D)
        D_I=np.zeros_like(D)
        img_ax = cv2.drawFrameAxes(img_undis, K,D_I, rvec, tvec, 5,thickness=2)
        img_out = img_ax
    else:
        img_ax = cv2.drawFrameAxes(im_with_aruco_board, K, D, rvec, tvec, 2.5,thickness=2)
        img_out = img_ax.copy()
    if debug:
        showImg(img_out)
    
    R, _ = cv2.Rodrigues(rvec)
    P = np.concatenate([R,tvec],axis=1)
    P = np.vstack([P,[[0,0,0,1]]])
    return P, img_points, obj_points,ids,img_out

def sort_1st_col(arr):
    sorted_indices = np.argsort(arr[:, 0])
    sorted_arr = arr[sorted_indices]
    return sorted_arr

def save_obs_img_pts(output_path,img_points_a,obj_points_a,ids_a):
    
        # pdb.set_trace()
        img_points_a = img_points_a[::4]
        obj_points_a = obj_points_a[::4]
        
        img_points_a = img_points_a.reshape(img_points_a.shape[0],2)
        obj_points_a = obj_points_a.reshape(obj_points_a.shape[0],3)
        
        img_points_a = np.hstack([ids_a,img_points_a])
        obj_points_a = np.hstack([ids_a,obj_points_a])
        #sort by ids
        img_points_a = sort_1st_col(img_points_a)
        obj_points_a = sort_1st_col(obj_points_a)
        
        np.savetxt(output_path+"/img_points_target.txt",img_points_a,fmt='%.8f')
        np.savetxt(output_path+"/obj_points_target.txt",obj_points_a,fmt='%.8f')


if __name__=="__main__":
    i_p=sys.argv[1]
    qtruck_name=sys.argv[2]
    cam_type=sys.argv[3]

    # dictionary = cv2.aruco.getPredefinedDictionary(dict=cv2.aruco.DICT_4X4_250)
    # board = cv2.aruco.CharucoBoard((4,3),
    #                             squareLength=0.420,
    #                             markerLength=0.315,
    #                             dictionary=dictionary,
    #                             ids=np.array([[82],[83],[84],[85],[86],[87]])
    # )
    # new
    ID_min,ID_max=187,187+7
    ids_list=[]
    for i in range(ID_min,ID_max,1):
        ids_list.append(i)
        print("ids: \n",ids_list)
    print(len(ids_list))
    dictionary = cv2.aruco.getPredefinedDictionary(dict=cv2.aruco.DICT_4X4_250)
    board = cv2.aruco.CharucoBoard((5,3),
                                squareLength=1.,
                                markerLength=0.750,
                                dictionary=dictionary,
                                ids=np.array(ids_list)
    )



    # code 
    K_dir = "/home/westwell/welldriver/wellpilot_config/calib/cameras_in_calib_file"
    K,D=getKD(osp.join(K_dir,qtruck_name,cam_type+".yaml"))
    image = cv2.imread(i_p)
    P, img_points, obj_points,ids,img_out = get_Eb2c(K,D,image,board,dictionary)
    np.savetxt("Eb2c.txt",P)

    cv2.imwrite("img_corners_"+cam_type+".jpg",img_out)
    print("get_Eb2c : \n",P)
    print("img_corners_"+cam_type+".jpg")